// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/serializers/presto_serializer.h>
#include <boost/random/uniform_int_distribution.hpp>
#include <melon/random.h>
#include <gtest/gtest.h>
#include <vector>
#include "melon/experimental/event_count.h"
#include <pollux/testing/gtest_utils.h>
#include <pollux/common/memory/byte_stream.h>
#include <pollux/functions/prestosql/types/ip_address_type.h>
#include <pollux/functions/prestosql/types/ip_prefix_type.h>
#include <pollux/functions/prestosql/types/timestamp_with_time_zone_type.h>
#include <pollux/serializers/presto_batch_vector_serializer.h>
#include <pollux/serializers/presto_vector_lexer.h>
#include <pollux/vector/vector_fuzzer.h>
#include <pollux/testing/vector/vector_test_base.h>

using namespace kumo::pollux;
using namespace kumo::pollux::test;

struct SerializeStats {
    int64_t actualSize{0};
    int64_t estimatedSize{0};
};

class Foo {
public:
    explicit Foo(int64_t id) : id_(id) {
    }

    int64_t id() const {
        return id_;
    }

    static std::shared_ptr<Foo> create(int64_t id) {
        // Return the same instance if the id is already in the map, to make
        // the operator== work with the instance before and after serde.
        if (instances_.find(id) != instances_.end()) {
            return instances_[id];
        }
        instances_[id] = std::make_shared<Foo>(id);
        return instances_[id];
    }

    static std::string serialize(const std::shared_ptr<Foo> &foo) {
        return std::to_string(foo->id_);
    }

    static std::shared_ptr<Foo> deserialize(const std::string &serialized) {
        if (serialized == "") {
            return nullptr;
        }
        return create(std::stoi(serialized));
    }

private:
    int64_t id_;
    static melon::F14FastMap<int64_t, std::shared_ptr<Foo> > instances_;
};

melon::F14FastMap<int64_t, std::shared_ptr<Foo> > Foo::instances_;

class PrestoSerializerTest
        : public ::testing::TestWithParam<common::CompressionKind>,
          public VectorBuilder {
protected:
    static void SetUpTestSuite() {
        if (!isRegisteredVectorSerde()) {
            serializer::presto::PrestoVectorSerde::registerVectorSerde();
        }
        memory::MemoryManager::testingSetInstance({});
    }

    void SetUp() override {
        serde_ = std::make_unique<serializer::presto::PrestoVectorSerde>();
    }

    void TearDown() override {
        OpaqueType::clearSerializationRegistry();
    }

    void sanityCheckEstimateSerializedSize(const RowVectorPtr &row_vector) {
        Scratch scratch;
        const auto numRows = row_vector->size();

        std::vector<IndexRange> rows(numRows);
        for (int i = 0; i < numRows; i++) {
            rows[i] = IndexRange{i, 1};
        }

        std::vector<vector_size_t> rowSizes(numRows, 0);
        std::vector<vector_size_t *> rawRowSizes(numRows);
        for (auto i = 0; i < numRows; i++) {
            rawRowSizes[i] = &rowSizes[i];
        }
        serde_->estimateSerializedSize(
            row_vector.get(),
            melon::Range(rows.data(), numRows),
            rawRowSizes.data(),
            scratch);
    }

    serializer::presto::PrestoVectorSerde::PrestoOptions getParamSerdeOptions(
        const serializer::presto::PrestoVectorSerde::PrestoOptions *
        serdeOptions) {
        const bool useLosslessTimestamp =
                serdeOptions == nullptr ? false : serdeOptions->useLosslessTimestamp;
        common::CompressionKind kind = GetParam();
        const bool nullsFirst =
                serdeOptions == nullptr ? false : serdeOptions->nullsFirst;
        const bool preserveEncodings =
                serdeOptions == nullptr ? false : serdeOptions->preserveEncodings;
        serializer::presto::PrestoVectorSerde::PrestoOptions paramOptions{
            useLosslessTimestamp, kind, 0.8, nullsFirst, preserveEncodings
        };

        return paramOptions;
    }

    SerializeStats serialize(
        const RowVectorPtr &row_vector,
        std::ostream *output,
        const serializer::presto::PrestoVectorSerde::PrestoOptions *serdeOptions,
        std::optional<melon::Range<const IndexRange *> > indexRanges = std::nullopt,
        std::optional<melon::Range<const vector_size_t *> > rows = std::nullopt,
        std::unique_ptr<IterativeVectorSerializer> *reuseSerializer = nullptr,
        std::unique_ptr<StreamArena> *reuseArena = nullptr) {
        auto streamInitialSize = output->tellp();
        sanityCheckEstimateSerializedSize(row_vector);

        std::unique_ptr<StreamArena> arena;
        auto rowType = as_row_type(row_vector->type());
        auto numRows = row_vector->size();
        auto paramOptions = getParamSerdeOptions(serdeOptions);
        std::unique_ptr<IterativeVectorSerializer> serializer;
        if (reuseSerializer && *reuseSerializer) {
            arena = std::move(*reuseArena);
            serializer = std::move(*reuseSerializer);
            serializer->clear();
        } else {
            arena = std::make_unique<StreamArena>(pool_.get());
            serializer = serde_->createIterativeSerializer(
                rowType, numRows, arena.get(), &paramOptions);
        }
        vector_size_t sizeEstimate = 0;

        Scratch scratch;
        if (indexRanges.has_value()) {
            raw_vector<vector_size_t *> sizes(indexRanges.value().size());
            std::fill(sizes.begin(), sizes.end(), &sizeEstimate);
            serde_->estimateSerializedSize(
                row_vector.get(), indexRanges.value(), sizes.data(), scratch);
            serializer->append(row_vector, indexRanges.value(), scratch);
        } else if (rows.has_value()) {
            raw_vector<vector_size_t> sizes(rows.value().size());
            std::vector<vector_size_t *> sizePointers(rows.value().size());
            for (vector_size_t i = 0; i < sizes.size(); ++i) {
                sizePointers[i] = &sizes[i];
            }
            serde_->estimateSerializedSize(
                row_vector.get(), rows.value(), sizePointers.data(), scratch);
            serializer->append(row_vector, rows.value(), scratch);
        } else {
            vector_size_t *sizes = &sizeEstimate;
            IndexRange range{0, row_vector->size()};
            serde_->estimateSerializedSize(
                row_vector.get(),
                melon::Range<const IndexRange *>(&range, 1),
                &sizes,
                scratch);
            serializer->append(row_vector);
        }
        auto size = serializer->maxSerializedSize();
        auto estimatePct = (100.0 * sizeEstimate) / static_cast<float>(size + 1);
        KLOG(INFO) << "Size=" << size << " estimate=" << sizeEstimate << " "
              << estimatePct << "%";
        kumo::pollux::serializer::presto::PrestoOutputStreamListener listener;
        OStreamOutputStream out(output, &listener);
        serializer->flush(&out);
        if (paramOptions.compressionKind == common::CompressionKind_NONE) {
            EXPECT_EQ(size, out.tellp() - streamInitialSize);
        } else {
            EXPECT_GE(size, out.tellp() - streamInitialSize);
        }
        if (reuseSerializer) {
            *reuseArena = std::move(arena);
            *reuseSerializer = std::move(serializer);
        }
        return {static_cast<int64_t>(size), sizeEstimate};
    }

    std::unique_ptr<ByteInputStream> toByteStream(const std::string &input) {
        ByteRange byteRange{
            reinterpret_cast<uint8_t *>(const_cast<char *>(input.data())),
            (int32_t) input.length(),
            0
        };
        return std::make_unique<BufferInputStream>(
            std::vector<ByteRange>{{byteRange}});
    }

    void validateLexer(
        const std::string &input,
        const serializer::presto::PrestoVectorSerde::PrestoOptions &
        paramOptions) {
        if (paramOptions.useLosslessTimestamp ||
            paramOptions.compressionKind !=
            common::CompressionKind::CompressionKind_NONE ||
            paramOptions.nullsFirst) {
            // Unsupported options
            return;
        }
        std::vector<serializer::presto::Token> tokens;
        const auto status = serializer::presto::PrestoVectorSerde::lex(
            input, tokens, &paramOptions);
        EXPECT_TRUE(status.ok()) << status.message();

        size_t tokenLengthSum = 0;
        for (size_t i = 0; i < tokens.size(); ++i) {
            const auto &token = tokens[i];

            // The lexer should not produce empty tokens
            EXPECT_NE(token.length, 0);
            // Compute tokenLengthSum to validate with input.size().
            tokenLengthSum += token.length;

            if (i > 0) {
                // The lexer should merge consecutive tokens of the same type
                EXPECT_NE(token.tokenType, tokens[i - 1].tokenType);
            }
        }
        EXPECT_EQ(tokenLengthSum, input.size());
    }

    RowVectorPtr deserialize(
        const RowTypePtr &rowType,
        const std::string &input,
        const serializer::presto::PrestoVectorSerde::PrestoOptions *serdeOptions,
        bool skipLexer = false) {
        auto byteStream = toByteStream(input);
        auto paramOptions = getParamSerdeOptions(serdeOptions);
        if (!skipLexer) {
            validateLexer(input, paramOptions);
        }
        RowVectorPtr result;
        serde_->deserialize(
            byteStream.get(), pool_.get(), rowType, &result, 0, &paramOptions);
        return result;
    }

    RowVectorPtr makeTestVector(vector_size_t size) {
        auto a =
                make_flat_vector<int64_t>(size, [](vector_size_t row) { return row; });
        auto b = make_flat_vector<double>(
            size, [](vector_size_t row) { return row * 0.1; });
        auto c = make_flat_vector<std::string>(size, [](vector_size_t row) {
            return row % 2 == 0 ? "LaaaaaaaaargeString" : "inlineStr";
        });

        std::vector<VectorPtr> childVectors = {a, b, c};

        return make_row_vector(childVectors);
    }

    RowVectorPtr wrapChildren(const RowVectorPtr &row) {
        auto children = row->children();
        std::vector<VectorPtr> newChildren = children;
        auto indices = make_indices(row->size(), [](auto row) { return row; });
        for (auto &child: newChildren) {
            child = BaseVector::wrap_in_dictionary(
                BufferPtr(nullptr), indices, row->size(), child);
        }
        return make_row_vector(newChildren);
    }

    void testRoundTrip(
        VectorPtr vector,
        const serializer::presto::PrestoVectorSerde::PrestoOptions *serdeOptions =
                nullptr) {
        std::unique_ptr<IterativeVectorSerializer> reuseSerializer;
        std::unique_ptr<StreamArena> reuseArena;
        auto row_vector = make_row_vector({vector});
        std::ostringstream out;
        serialize(
            row_vector,
            &out,
            serdeOptions,
            std::nullopt,
            std::nullopt,
            &reuseSerializer,
            &reuseArena);

        auto rowType = as_row_type(row_vector->type());
        auto deserialized = deserialize(rowType, out.str(), serdeOptions);
        assertEqualVectors(deserialized, row_vector);

        if (row_vector->size() < 3) {
            return;
        }

        // Split input into 3 batches. Serialize each separately. Then, deserialize
        // all into one vector.
        auto splits = split(row_vector, 3);
        std::vector<std::string> serialized;
        for (const auto &split: splits) {
            std::ostringstream out;
            serialize(
                split,
                &out,
                serdeOptions,
                std::nullopt,
                std::nullopt,
                &reuseSerializer,
                &reuseArena);
            serialized.push_back(out.str());
        }

        auto paramOptions = getParamSerdeOptions(serdeOptions);
        RowVectorPtr result;
        vector_size_t offset = 0;
        for (auto i = 0; i < serialized.size(); ++i) {
            auto byteStream = toByteStream(serialized[i]);
            serde_->deserialize(
                byteStream.get(),
                pool_.get(),
                rowType,
                &result,
                offset,
                &paramOptions);
            offset = result->size();
        }

        assertEqualVectors(result, row_vector);

        // Serialize the vector with even and odd rows in different partitions.
        auto even =
                make_indices(row_vector->size() / 2, [&](auto row) { return row * 2; });
        auto odd = make_indices(
            (row_vector->size() - 1) / 2, [&](auto row) { return (row * 2) + 1; });
        testSerializeRows(
            row_vector, even, serdeOptions, &reuseSerializer, &reuseArena);
        auto oddStats = testSerializeRows(
            row_vector, odd, serdeOptions, &reuseSerializer, &reuseArena);
        auto wrappedRowVector = wrapChildren(row_vector);
        auto wrappedStats = testSerializeRows(
            wrappedRowVector, odd, serdeOptions, &reuseSerializer, &reuseArena);
        EXPECT_EQ(oddStats.estimatedSize, wrappedStats.estimatedSize);
        // The second serialization may come out smaller if encoding is better.
        EXPECT_GE(oddStats.actualSize, wrappedStats.actualSize);
    }

    void testLexer(
        VectorPtr vector,
        const serializer::presto::PrestoVectorSerde::PrestoOptions *serdeOptions =
                nullptr) {
        auto row_vector = make_row_vector({vector});
        std::ostringstream out;
        serialize(row_vector, &out, serdeOptions);
        validateLexer(out.str(), getParamSerdeOptions(serdeOptions));
    }

    SerializeStats testSerializeRows(
        const RowVectorPtr &row_vector,
        BufferPtr indices,
        const serializer::presto::PrestoVectorSerde::PrestoOptions *serdeOptions,
        std::unique_ptr<IterativeVectorSerializer> *reuseSerializer,
        std::unique_ptr<StreamArena> *reuseArena) {
        std::ostringstream out;
        auto rows = melon::Range<const vector_size_t *>(
            indices->as<vector_size_t>(), indices->size() / sizeof(vector_size_t));
        auto stats = serialize(
            row_vector,
            &out,
            serdeOptions,
            std::nullopt,
            rows,
            reuseSerializer,
            reuseArena);

        auto rowType = as_row_type(row_vector->type());
        auto deserialized = deserialize(rowType, out.str(), serdeOptions);
        assertEqualVectors(
            deserialized,
            BaseVector::wrap_in_dictionary(
                BufferPtr(nullptr),
                indices,
                indices->size() / sizeof(vector_size_t),
                row_vector));
        return stats;
    }

    void assertEqualEncoding(
        const RowVectorPtr &expected,
        const RowVectorPtr &actual,
        // If true, we allow the encodings of actual and expected to differ if a
        // dictionary was flattened.
        bool allowFlatteningDictionaries = false) {
        for (auto i = 0; i < expected->childrenSize(); ++i) {
            auto expectedEncoding = expected->childAt(i)->encoding();
            auto actualEncoding = actual->childAt(i)->encoding();

            if (allowFlatteningDictionaries &&
                actualEncoding == VectorEncoding::Simple::FLAT &&
                expectedEncoding == VectorEncoding::Simple::DICTIONARY) {
                continue;
            }

            POLLUX_CHECK_EQ(actualEncoding, expectedEncoding);

            if (expectedEncoding == VectorEncoding::Simple::ROW) {
                assertEqualEncoding(
                    std::dynamic_pointer_cast<RowVector>(expected->childAt(i)),
                    std::dynamic_pointer_cast<RowVector>(actual->childAt(i)),
                    allowFlatteningDictionaries);
            }
        }
    }

    void verifySerializedEncodedData(
        const RowVectorPtr &original,
        const std::string &serialized,
        bool allowFlatteningDictionaries,
        const serializer::presto::PrestoVectorSerde::PrestoOptions *
        serdeOptions) {
        auto rowType = as_row_type(original->type());
        auto deserialized = deserialize(rowType, serialized, serdeOptions);

        assertEqualVectors(original, deserialized);
        // Dictionaries may get flattened depending on the nature of the data.
        assertEqualEncoding(original, deserialized, allowFlatteningDictionaries);

        // Deserialize 3 times while appending to a single vector.
        auto paramOptions = getParamSerdeOptions(serdeOptions);
        RowVectorPtr result;
        vector_size_t offset = 0;
        for (auto i = 0; i < 3; ++i) {
            auto byteStream = toByteStream(serialized);
            serde_->deserialize(
                byteStream.get(),
                pool_.get(),
                rowType,
                &result,
                offset,
                &paramOptions);
            offset = result->size();
        }

        auto expected =
                BaseVector::create(original->type(), original->size() * 3, pool());
        for (auto i = 0; i < 3; ++i) {
            expected->copy(original.get(), original->size() * i, 0, original->size());
        }

        assertEqualVectors(expected, result);
    }

    void makePermutations(
        const std::vector<VectorPtr> &vectors,
        int32_t size,
        std::vector<VectorPtr> &items,
        std::vector<std::vector<VectorPtr> > &result) {
        if (size == items.size()) {
            result.push_back(items);
            return;
        }
        for (const auto &vector: vectors) {
            items.push_back(vector);
            makePermutations(vectors, size, items, result);
            items.pop_back();
        }
    }

    // tests combining encodings in serialization and deserialization. Serializes
    // each of 'vectors' with encoding, then reads them back into a single vector.
    void testEncodedConcatenation(
        std::vector<VectorPtr> &vectors,
        const serializer::presto::PrestoVectorSerde::PrestoOptions *serdeOptions =
                nullptr) {
        std::vector<std::string> pieces;
        std::vector<std::string> reusedPieces;
        auto rowType = ROW({{"f", vectors[0]->type()}});
        auto concatenation = BaseVector::create(rowType, 0, pool_.get());
        auto arena = std::make_unique<StreamArena>(pool_.get());
        auto paramOptions = getParamSerdeOptions(serdeOptions);
        auto serializer = serde_->createIterativeSerializer(
            rowType, 10, arena.get(), &paramOptions);
        auto reusedSerializer = serde_->createIterativeSerializer(
            rowType, 10, arena.get(), &paramOptions);

        kumo::pollux::serializer::presto::PrestoOutputStreamListener listener;
        std::ostringstream allOut;
        OStreamOutputStream allOutStream(&allOut, &listener);
        serializer->flush(&allOutStream);

        auto allDeserialized = deserialize(rowType, allOut.str(), &paramOptions);
        assertEqualVectors(allDeserialized, concatenation);
        RowVectorPtr deserialized =
                BaseVector::create<RowVector>(rowType, 0, pool_.get());
        for (auto pieceIdx = 0; pieceIdx < pieces.size(); ++pieceIdx) {
            auto piece = pieces[pieceIdx];
            auto byteStream = toByteStream(piece);
            serde_->deserialize(
                byteStream.get(),
                pool_.get(),
                rowType,
                &deserialized,
                deserialized->size(),
                &paramOptions);

            RowVectorPtr single =
                    BaseVector::create<RowVector>(rowType, 0, pool_.get());
            byteStream = toByteStream(piece);
            serde_->deserialize(
                byteStream.get(), pool_.get(), rowType, &single, 0, &paramOptions);
            assertEqualVectors(single->childAt(0), vectors[pieceIdx]);

            RowVectorPtr single2 =
                    BaseVector::create<RowVector>(rowType, 0, pool_.get());
            byteStream = toByteStream(reusedPieces[pieceIdx]);
            serde_->deserialize(
                byteStream.get(), pool_.get(), rowType, &single2, 0, &paramOptions);
            assertEqualVectors(single2->childAt(0), vectors[pieceIdx]);
        }
        assertEqualVectors(concatenation, deserialized);
    }

    int64_t randInt(int64_t min, int64_t max) {
        return min + melon::Random::rand64(rng_) % (max - min);
    }

    void randomOptions(VectorFuzzer::Options &options, bool isFirst) {
        options.nullRatio = randInt(1, 10) < 3 ? 0.0 : randInt(1, 10) / 10.0;
        options.stringLength = randInt(1, 100);
        options.stringVariableLength = true;
        options.containerLength = randInt(1, 50);
        options.containerVariableLength = true;
        options.complexElementsMaxSize = 20000;
        options.maxConstantContainerSize = 2;
        options.normalizeMapKeys = randInt(0, 20) < 16;
        if (isFirst) {
            options.timestampPrecision =
                    static_cast<VectorFuzzer::Options::TimestampPrecision>(randInt(0, 3));
        }
        options.allowLazyVector = false;
    }

    void serializeBatch(
        const RowVectorPtr &row_vector,
        std::ostream *output,
        const serializer::presto::PrestoVectorSerde::PrestoOptions *
        serdeOptions) {
        kumo::pollux::serializer::presto::PrestoOutputStreamListener listener;
        OStreamOutputStream out(output, &listener);
        auto paramOptions = getParamSerdeOptions(serdeOptions);

        auto serializer = serde_->createBatchSerializer(pool_.get(), &paramOptions);
        serializer->serialize(row_vector, &out);
    }

    void testBatchVectorSerializerRoundTrip(
        const RowVectorPtr &data,
        bool allowFlatteningDictionaries = false,
        const serializer::presto::PrestoVectorSerde::PrestoOptions *serdeOptions =
                nullptr) {
        std::ostringstream out;
        serializeBatch(data, &out, serdeOptions);
        const auto serialized = out.str();

        verifySerializedEncodedData(
            data, serialized, allowFlatteningDictionaries, serdeOptions);
    }

    RowVectorPtr encodingsTestVector() {
        // String is variable length, so this ensures the data isn't flattened.
        auto baseNoNulls = make_flat_vector<std::string>({"a", "b", "c", "d"});
        auto baseWithNulls =
                make_nullable_flat_vector<std::string>({"a", std::nullopt, "b", "c"});
        auto baseArray =
                make_array_vector<int32_t>({{1, 2, 3}, {}, {4, 5}, {6, 7, 8, 9, 10}});
        auto indices = make_indices(8, [](auto row) { return row / 2; });

        return make_row_vector({
            BaseVector::wrap_in_dictionary(nullptr, indices, 8, baseNoNulls),
            BaseVector::wrap_in_dictionary(nullptr, indices, 8, baseWithNulls),
            BaseVector::wrap_in_dictionary(nullptr, indices, 8, baseArray),
            BaseVector::create_constant(INTEGER(), 123, 8, pool_.get()),
            BaseVector::create_null_constant(VARCHAR(), 8, pool_.get()),
            BaseVector::wrap_in_constant(8, 1, baseArray),
            BaseVector::wrap_in_constant(8, 2, baseArray),
            make_row_vector({
                BaseVector::wrap_in_dictionary(nullptr, indices, 8, baseNoNulls),
                BaseVector::wrap_in_dictionary(nullptr, indices, 8, baseWithNulls),
                BaseVector::wrap_in_dictionary(nullptr, indices, 8, baseArray),
                BaseVector::create_constant(INTEGER(), 123, 8, pool_.get()),
                BaseVector::create_null_constant(VARCHAR(), 8, pool_.get()),
                BaseVector::wrap_in_constant(8, 1, baseArray),
                BaseVector::wrap_in_constant(8, 2, baseArray),
                make_row_vector({
                    BaseVector::wrap_in_dictionary(
                        nullptr, indices, 8, baseWithNulls),
                    BaseVector::create_constant(INTEGER(), 123, 8, pool_.get()),
                    BaseVector::wrap_in_constant(8, 2, baseArray),
                }),
            }),
        });
    }

    RowVectorPtr encodingsArrayElementsTestVector() {
        auto baseNoNulls = make_flat_vector<int64_t>({1, 2, 3, 4});
        auto baseWithNulls =
                make_nullable_flat_vector<int32_t>({1, std::nullopt, 2, 3});
        auto baseArray =
                make_array_vector<int32_t>({{1, 2, 3}, {}, {4, 5}, {6, 7, 8, 9, 10}});
        auto elementIndices = make_indices(16, [](auto row) { return row / 4; });
        std::vector<vector_size_t> offsets{0, 2, 4, 6, 8, 10, 12, 14, 16};

        return make_row_vector({
            make_array_vector(
                offsets,
                BaseVector::wrap_in_dictionary(
                    nullptr, elementIndices, 16, baseNoNulls)),
            make_array_vector(
                offsets,
                BaseVector::wrap_in_dictionary(
                    nullptr, elementIndices, 16, baseWithNulls)),
            make_array_vector(
                offsets,
                BaseVector::wrap_in_dictionary(
                    nullptr, elementIndices, 16, baseArray)),
            make_array_vector(
                offsets,
                BaseVector::create_constant(INTEGER(), 123, 16, pool_.get())),
            make_array_vector(
                offsets,
                BaseVector::create_null_constant(VARCHAR(), 16, pool_.get())),
            make_array_vector(offsets, BaseVector::wrap_in_constant(16, 1, baseArray)),
            make_row_vector({
                make_array_vector(
                    offsets,
                    BaseVector::wrap_in_dictionary(
                        nullptr, elementIndices, 16, baseNoNulls)),
                make_array_vector(
                    offsets,
                    BaseVector::wrap_in_dictionary(
                        nullptr, elementIndices, 16, baseWithNulls)),
                make_array_vector(
                    offsets,
                    BaseVector::wrap_in_dictionary(
                        nullptr, elementIndices, 16, baseArray)),
                make_array_vector(
                    offsets,
                    BaseVector::create_constant(INTEGER(), 123, 16, pool_.get())),
                make_array_vector(
                    offsets,
                    BaseVector::create_null_constant(VARCHAR(), 16, pool_.get())),
                make_array_vector(
                    offsets, BaseVector::wrap_in_constant(16, 1, baseArray)),
            }),
        });
    }

    RowVectorPtr encodingsMapValuesTestVector() {
        auto baseNoNulls = make_flat_vector<int64_t>({1, 2, 3, 4});
        auto baseWithNulls =
                make_nullable_flat_vector<int32_t>({1, std::nullopt, 2, 3});
        auto baseArray =
                make_array_vector<int32_t>({{1, 2, 3}, {}, {4, 5}, {6, 7, 8, 9, 10}});
        auto valueIndices = make_indices(16, [](auto row) { return row / 4; });
        std::vector<vector_size_t> offsets{0, 2, 4, 6, 8, 10, 12, 14, 16};
        auto mapKeys = make_flat_vector<int32_t>(16, [](auto row) { return row; });

        return make_row_vector({
            make_map_vector(
                offsets,
                mapKeys,
                BaseVector::wrap_in_dictionary(
                    nullptr, valueIndices, 16, baseNoNulls)),
            make_map_vector(
                offsets,
                mapKeys,
                BaseVector::wrap_in_dictionary(
                    nullptr, valueIndices, 16, baseWithNulls)),
            make_map_vector(
                offsets,
                mapKeys,
                BaseVector::wrap_in_dictionary(nullptr, valueIndices, 16, baseArray)),
            make_map_vector(
                offsets,
                mapKeys,
                BaseVector::create_constant(INTEGER(), 123, 16, pool_.get())),
            make_map_vector(
                offsets,
                mapKeys,
                BaseVector::create_null_constant(VARCHAR(), 16, pool_.get())),
            make_map_vector(
                offsets, mapKeys, BaseVector::wrap_in_constant(16, 1, baseArray)),
            make_row_vector({
                make_map_vector(
                    offsets,
                    mapKeys,
                    BaseVector::wrap_in_dictionary(
                        nullptr, valueIndices, 16, baseNoNulls)),
                make_map_vector(
                    offsets,
                    mapKeys,
                    BaseVector::wrap_in_dictionary(
                        nullptr, valueIndices, 16, baseWithNulls)),
                make_map_vector(
                    offsets,
                    mapKeys,
                    BaseVector::wrap_in_dictionary(
                        nullptr, valueIndices, 16, baseArray)),
                make_map_vector(
                    offsets,
                    mapKeys,
                    BaseVector::create_constant(INTEGER(), 123, 16, pool_.get())),
                make_map_vector(
                    offsets,
                    mapKeys,
                    BaseVector::create_null_constant(VARCHAR(), 16, pool_.get())),
                make_map_vector(
                    offsets, mapKeys, BaseVector::wrap_in_constant(16, 1, baseArray)),
            }),
        });
    }

    void testMinimalDictionaryEncoding(
        vector_size_t numRows,
        vector_size_t alphabetSize) {
        // This factor is used to ensure we have some repetition in the dictionary
        // in each of the cases to ensure the serializer doesn't flatten the data.
        auto factor = alphabetSize <= numRows ? 2 : alphabetSize / numRows * 2;
        // String is variable length, so this ensures the data isn't flattened.
        auto base = make_flat_vector<std::string>(
            alphabetSize, [](vector_size_t row) { return fmt::format("{}", row); });
        auto evenIndices = make_indices(numRows, [alphabetSize, factor](auto row) {
            return (row * factor) % alphabetSize;
        });
        auto oddIndices = make_indices(numRows, [alphabetSize, factor](auto row) {
            return (row * factor + 1) % alphabetSize;
        });
        auto prefixIndices = make_indices(numRows, [alphabetSize, factor](auto row) {
            return row % (alphabetSize / factor);
        });
        auto suffixIndices = make_indices(numRows, [alphabetSize, factor](auto row) {
            return row % (alphabetSize / factor) +
                   (factor - 1) * (alphabetSize / factor);
        });

        auto rows = make_row_vector({
            BaseVector::wrap_in_dictionary(nullptr, evenIndices, numRows, base),
            BaseVector::wrap_in_dictionary(nullptr, oddIndices, numRows, base),
            BaseVector::wrap_in_dictionary(nullptr, prefixIndices, numRows, base),
            BaseVector::wrap_in_dictionary(nullptr, suffixIndices, numRows, base),
        });

        std::ostringstream out;
        serializeBatch(rows, &out, /*serdeOptions=*/nullptr);
        const auto serialized = out.str();

        auto rowType = as_row_type(rows->type());
        auto deserialized =
                deserialize(rowType, serialized, /*serdeOptions=*/nullptr);

        assertEqualVectors(rows, deserialized);
        assertEqualEncoding(rows, deserialized);

        ASSERT_EQ(
            deserialized->childAt(0)
            ->as<DictionaryVector<StringView>>()
            ->value_vector()
            ->size(),
            alphabetSize / factor);
        ASSERT_EQ(
            deserialized->childAt(1)
            ->as<DictionaryVector<StringView>>()
            ->value_vector()
            ->size(),
            alphabetSize / factor);
        ASSERT_EQ(
            deserialized->childAt(2)
            ->as<DictionaryVector<StringView>>()
            ->value_vector()
            ->size(),
            alphabetSize / factor);
        ASSERT_EQ(
            deserialized->childAt(3)
            ->as<DictionaryVector<StringView>>()
            ->value_vector()
            ->size(),
            alphabetSize / factor);
    }

    RowVectorPtr makeEmptyTestVector() {
        return make_row_vector(
            {"a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l"},
            {
                make_flat_vector<bool>({}),
                make_flat_vector<int8_t>({}),
                make_flat_vector<int16_t>({}),
                make_flat_vector<int32_t>({}),
                make_flat_vector<int64_t>({}),
                make_flat_vector<float>({}),
                make_flat_vector<double>({}),
                make_flat_vector<StringView>({}),
                make_flat_vector<Timestamp>({}),
                make_row_vector(
                    {"1", "2"},
                    {make_flat_vector<StringView>({}), make_flat_vector<int32_t>({})}),
                make_array_vector<int32_t>({}),
                make_map_vector<StringView, int32_t>({})
            });
    }

    void testDeserializeSingleColumn(
        const std::string &serializedData,
        const VectorPtr &expected) {
        auto byteStream = toByteStream(serializedData);
        VectorPtr deserialized;
        serde_->deserializeSingleColumn(
            byteStream.get(), pool(), expected->type(), &deserialized, nullptr);
        assertEqualVectors(expected, deserialized);
    }

    std::unique_ptr<serializer::presto::PrestoVectorSerde> serde_;
    melon::Random::DefaultGenerator rng_;
};

TEST_P(PrestoSerializerTest, basic) {
    vector_size_t numRows = 1'000;
    auto row_vector = makeTestVector(numRows);
    testRoundTrip(row_vector);
}

TEST_P(PrestoSerializerTest, basicLarge) {
#ifndef NDEBUG
    constexpr vector_size_t kNumRows = 20'000;
#else
  constexpr vector_size_t kNumRows = 200'000;
#endif
    auto row_vector = make_row_vector({
        make_flat_vector<int64_t>(kNumRows, [](vector_size_t row) { return row; }),
        make_flat_vector<std::string>(
            kNumRows, [](vector_size_t) { return std::string(2048, 'x'); }),
    });
    testRoundTrip(std::move(row_vector));
}

/// Test serialization of a dictionary vector that adds nulls to the base
/// vector.
TEST_P(PrestoSerializerTest, dictionaryWithExtraNulls) {
    vector_size_t size = 1'000;

    auto base = make_flat_vector<int64_t>(10, [](auto row) { return row; });

    BufferPtr nulls = AlignedBuffer::allocate<bool>(size, pool_.get());
    auto rawNulls = nulls->asMutable<uint64_t>();
    for (auto i = 0; i < size; i++) {
        bits::set_null(rawNulls, i, i % 5 == 0);
    }

    BufferPtr indices = AlignedBuffer::allocate<vector_size_t>(size, pool_.get());
    auto rawIndices = indices->asMutable<vector_size_t>();
    for (auto i = 0; i < size; i++) {
        if (i % 5 != 0) {
            rawIndices[i] = i % 10;
        }
    }

    auto dictionary = BaseVector::wrap_in_dictionary(nulls, indices, size, base);
    testRoundTrip(dictionary);
}

TEST_P(PrestoSerializerTest, emptyPage) {
    auto row_vector = makeEmptyTestVector();

    std::ostringstream out;
    serialize(row_vector, &out, nullptr);

    auto rowType = as_row_type(row_vector->type());
    auto deserialized = deserialize(rowType, out.str(), nullptr);
    assertEqualVectors(deserialized, row_vector);
}

TEST_P(PrestoSerializerTest, invalidPage) {
    auto row_vector = makeEmptyTestVector();

    std::ostringstream out;
    serialize(row_vector, &out, nullptr);

    auto invalidPage = ""; // empty string
    auto rowType = as_row_type(row_vector->type());
    POLLUX_ASSERT_THROW(
        deserialize(rowType, invalidPage, nullptr, true /*skipLexer*/),
        "PrestoPage header is invalid: 0 bytes for header");
}

TEST_P(PrestoSerializerTest, initMemory) {
    const auto numRows = 100;
    auto testFunc = [&](TypePtr type, int64_t expectedBytes) {
        const auto poolMemUsage = pool_->usedBytes();
        auto arena = std::make_unique<StreamArena>(pool_.get());
        const auto paramOptions = getParamSerdeOptions(nullptr);
        const auto rowType = ROW({type});
        const auto serializer = serde_->createIterativeSerializer(
            rowType, numRows, arena.get(), &paramOptions);
        ASSERT_EQ(pool_->usedBytes() - poolMemUsage, expectedBytes);
    };

    // 8192 for slab allocation for 64-bit header + 384 for memory pool allocated
    // std structures
    testFunc(BOOLEAN(), 8576);
    testFunc(TINYINT(), 8576);
    testFunc(SMALLINT(), 8576);
    testFunc(INTEGER(), 8576);
    testFunc(BIGINT(), 8576);
    testFunc(REAL(), 8576);
    testFunc(DOUBLE(), 8576);
    testFunc(VARCHAR(), 8576);
    testFunc(TIMESTAMP(), 8576);
    // For nested types, 2 pages allocation quantum for first offset (0).
    testFunc(ROW({VARCHAR()}), 8960);
    testFunc(ARRAY(INTEGER()), 8960);
    testFunc(MAP(VARCHAR(), INTEGER()), 9280);
}

TEST_P(PrestoSerializerTest, serializeNoRowsSelected) {
    std::ostringstream out;
    kumo::pollux::serializer::presto::PrestoOutputStreamListener listener;
    OStreamOutputStream output(&out, &listener);
    auto paramOptions = getParamSerdeOptions(nullptr);
    auto arena = std::make_unique<StreamArena>(pool_.get());
    auto row_vector = make_row_vector(
        {"a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l"},
        {
            make_flat_vector<bool>({true, false, true, false}),
            make_flat_vector<int8_t>({1, 2, 3, 4}),
            make_flat_vector<int16_t>({5, 6, 7, 8}),
            make_flat_vector<int32_t>({9, 10, 11, 12}),
            make_flat_vector<int64_t>({13, 14, 15, 16}),
            make_flat_vector<float>({17.0, 18.0, 19.0, 20.0}),
            make_flat_vector<double>({21.0, 22.0, 23.0, 24.0}),
            make_flat_vector<StringView>({"25", "26", "27", "28"}),
            make_flat_vector<Timestamp>({{29, 30}, {32, 32}, {33, 34}, {35, 36}}),
            make_row_vector(
                {"1", "2"},
                {
                    make_flat_vector<StringView>({"37", "38", "39", "40"}),
                    make_flat_vector<int32_t>({41, 42, 43, 44})
                }),
            make_array_vector<int32_t>({{45, 46}, {47, 48}, {49, 50}, {51, 52}}),
            make_map_vector<StringView, int32_t>(
                {
                    {{"53", 54}, {"55", 56}},
                    {{"57", 58}, {"59", 60}},
                    {{"61", 62}, {"63", 64}},
                    {{"65", 66}, {"67", 68}}
                })
        });
    const IndexRange noRows{0, 0};

    auto serializer = serde_->createIterativeSerializer(
        as_row_type(row_vector->type()),
        row_vector->size(),
        arena.get(),
        &paramOptions);

    serializer->append(row_vector, melon::Range(&noRows, 1));
    serializer->flush(&output);

    auto expected = makeEmptyTestVector();
    auto rowType = as_row_type(expected->type());
    auto deserialized = deserialize(rowType, out.str(), nullptr);
    assertEqualVectors(deserialized, expected);
}

TEST_P(PrestoSerializerTest, emptyArray) {
    auto array_vector = make_array_vector<int32_t>(
        1'000,
        [](vector_size_t row) { return row % 5; },
        [](vector_size_t row) { return row; });

    testRoundTrip(array_vector);
}

TEST_P(PrestoSerializerTest, emptyMap) {
    auto map_vector = make_map_vector<int32_t, int32_t>(
        1'000,
        [](vector_size_t row) { return row % 5; },
        [](vector_size_t row) { return row; },
        [](vector_size_t row) { return row * 2; });

    testRoundTrip(map_vector);
}

TEST_P(PrestoSerializerTest, timestampWithTimeZone) {
    auto timestamp = make_flat_vector<int64_t>(
        100,
        [](auto row) { return pack(10'000 + row, row % 37); },
        /* is_null_at */ nullptr,
        TIMESTAMP_WITH_TIME_ZONE());

    testRoundTrip(timestamp);

    // Add some nulls.
    for (auto i = 0; i < 100; i += 7) {
        timestamp->set_null(i, true);
    }
    testRoundTrip(timestamp);
}

TEST_P(PrestoSerializerTest, intervalDayTime) {
    auto vector = make_flat_vector<int64_t>(
        100,
        [](auto row) { return row + melon::Random::rand32(); },
        nullptr, // nullAt
        INTERVAL_DAY_TIME());

    testRoundTrip(vector);

    // Add some nulls.
    for (auto i = 0; i < 100; i += 7) {
        vector->set_null(i, true);
    }
    testRoundTrip(vector);
}

TEST_P(PrestoSerializerTest, unknown) {
    // Verify vectors of UNKNOWN type. Also verifies a special case where a
    // vector, not of UNKNOWN type and with all nulls is serialized as an UNKNOWN
    // type having BYTE_ARRAY encoding.
    auto testAllNullSerializedAsUnknown = [&](VectorPtr vector,
                                              TypePtr outputType) {
        auto row_vector = make_row_vector({vector});
        auto expected = make_row_vector(
            {BaseVector::create_null_constant(outputType, vector->size(), pool())});
        std::ostringstream out;
        serialize(row_vector, &out, nullptr);

        auto rowType = as_row_type(expected->type());
        auto deserialized = deserialize(rowType, out.str(), nullptr);
        assertEqualVectors(expected, deserialized);

        if (row_vector->size() < 3) {
            return;
        }

        // Split input into 3 batches. Serialize each separately. Then, deserialize
        // all into one vector.
        auto splits = split(row_vector, 3);
        std::vector<std::string> serialized;
        for (const auto &split: splits) {
            std::ostringstream oss;
            serialize(split, &oss, nullptr);
            serialized.push_back(oss.str());
        }

        auto paramOptions = getParamSerdeOptions(nullptr);
        RowVectorPtr result;
        vector_size_t offset = 0;
        for (auto i = 0; i < serialized.size(); ++i) {
            auto byteStream = toByteStream(serialized[i]);
            serde_->deserialize(
                byteStream.get(),
                pool_.get(),
                rowType,
                &result,
                offset,
                &paramOptions);
            offset = result->size();
        }

        assertEqualVectors(expected, result);
    };

    const vector_size_t size = 123;
    auto constant_vector =
            BaseVector::create_null_constant(UNKNOWN(), size, pool_.get());
    testRoundTrip(constant_vector);
    testAllNullSerializedAsUnknown(constant_vector, BIGINT());

    auto flat_vector = BaseVector::create(UNKNOWN(), size, pool_.get());
    for (auto i = 0; i < size; i++) {
        flat_vector->set_null(i, true);
    }
    testRoundTrip(flat_vector);
    testAllNullSerializedAsUnknown(flat_vector, BIGINT());
}

TEST_P(PrestoSerializerTest, multiPage) {
    std::ostringstream out;
    std::vector<RowVectorPtr> testVectors;
    // Note: Page of size 1250 is a slight increase in size that initiates string
    // buffer re-use.
    for (int size: {1234, 1250, 538, 2408}) {
        auto vec = makeTestVector(size);
        serialize(vec, &out, nullptr);
        testVectors.push_back(std::move(vec));
    }

    auto bytes = out.str();

    auto rowType = as_row_type(testVectors[0]->type());
    auto byteStream = toByteStream(bytes);

    RowVectorPtr deserialized;
    auto paramOptions = getParamSerdeOptions(nullptr);

    for (int i = 0; i < testVectors.size(); i++) {
        RowVectorPtr &vec = testVectors[i];
        serde_->deserialize(
            byteStream.get(),
            pool_.get(),
            rowType,
            &deserialized,
            0,
            &paramOptions);
        if (i < testVectors.size() - 1) {
            ASSERT_FALSE(byteStream->atEnd());
        } else {
            ASSERT_TRUE(byteStream->atEnd());
        }
        assertEqualVectors(deserialized, vec);
        deserialized->validate({});
    }
}

TEST_P(PrestoSerializerTest, timestampWithNanosecondPrecision) {
    // Verify that nanosecond precision is preserved when the right options are
    // passed to the serde.
    const serializer::presto::PrestoVectorSerde::PrestoOptions
            kUseLosslessTimestampOptions(
                true, common::CompressionKind::CompressionKind_NONE);
    auto timestamp = make_flat_vector<Timestamp>(
        {
            Timestamp{0, 0},
            Timestamp{12, 0},
            Timestamp{0, 17'123'456},
            Timestamp{1, 17'123'456},
            Timestamp{-1, 17'123'456}
        });
    testRoundTrip(timestamp, &kUseLosslessTimestampOptions);

    // Verify that precision is lost when no option is passed to the serde.
    auto timestampMillis = make_flat_vector<Timestamp>(
        {
            Timestamp{0, 0},
            Timestamp{12, 0},
            Timestamp{0, 17'000'000},
            Timestamp{1, 17'000'000},
            Timestamp{-1, 17'000'000}
        });
    auto inputRowVector = make_row_vector({timestamp});
    auto expectedOutputWithLostPrecision = make_row_vector({timestampMillis});
    std::ostringstream out;
    serialize(inputRowVector, &out, {});
    auto rowType = as_row_type(inputRowVector->type());
    auto deserialized = deserialize(rowType, out.str(), {});
    assertEqualVectors(deserialized, expectedOutputWithLostPrecision);
}

TEST_P(PrestoSerializerTest, longDecimal) {
    std::vector<int128_t> decimalValues(102);
    decimalValues[0] = DecimalUtil::kLongDecimalMin;
    for (int row = 1; row < 101; row++) {
        decimalValues[row] = row - 50;
    }
    decimalValues[101] = DecimalUtil::kLongDecimalMax;
    auto vector = make_flat_vector<int128_t>(decimalValues, DECIMAL(20, 5));

    testRoundTrip(vector);

    // Add some nulls.
    for (auto i = 0; i < 102; i += 7) {
        vector->set_null(i, true);
    }
    testRoundTrip(vector);
}

TEST_P(PrestoSerializerTest, ipprefix) {
    std::vector<int128_t> randomIpAddresss;
    randomIpAddresss.reserve(100);
    for (auto i = 0; i < 100; i++) {
        randomIpAddresss.push_back(
            HugeInt::build(melon::Random::rand64(), melon::Random::rand64()));
    }
    auto ipaddress = make_flat_vector<int128_t>(
        100,
        [&](auto row) { return randomIpAddresss.at(row); },
        /* is_null_at */ nullptr,
        IPADDRESS());

    auto prefix = make_flat_vector<int8_t>(
        100,
        [&](auto row) {
            if (randomIpAddresss.at(row) >= 4294967296) {
                return melon::Random::rand32(33, 65);
            }
            return melon::Random::rand32(0, 33);
        },
        /* is_null_at */ nullptr,
        TINYINT());

    auto vector = std::make_shared<RowVector>(
        pool_.get(),
        IPPREFIX(),
        BufferPtr(nullptr),
        100,
        std::vector<VectorPtr>{ipaddress, prefix});

    testRoundTrip(vector);

    // Add some nulls.
    for (auto i = 0; i < 100; i += 7) {
        vector->set_null(i, true);
    }
    testRoundTrip(vector);

    // Test that if the vector is wrapped, we can still properly
    // deserialize the ipprefix type
    auto oneIndex = make_indices(100, [](auto row) { return row; });
    auto wrapped_vector = make_row_vector({
        BaseVector::wrap_in_dictionary(
            make_nulls(100, [](auto row) { return row % 2 == 0; }),
            oneIndex,
            100,
            vector),
    });
    testRoundTrip(wrapped_vector);
}

TEST_P(PrestoSerializerTest, ipaddress) {
    auto ipaddress = make_flat_vector<int128_t>(
        100,
        [](auto row) {
            return HugeInt::build(melon::Random::rand64(), melon::Random::rand64());
        },
        /* is_null_at */ nullptr,
        IPADDRESS());

    testRoundTrip(ipaddress);

    // Add some nulls.
    for (auto i = 0; i < 100; i += 7) {
        ipaddress->set_null(i, true);
    }
    testRoundTrip(ipaddress);
}

TEST_P(PrestoSerializerTest, uuid) {
    auto vector = make_flat_vector<int128_t>(
        200, [](vector_size_t row) { return (int128_t) 0xD1 << row % 120; });

    testRoundTrip(vector);

    // Add some nulls.
    for (auto i = 0; i < vector->size(); i += 7) {
        vector->set_null(i, true);
    }
    testRoundTrip(vector);
}

// Test that hierarchically encoded columns (rows) have their encodings
// preserved by the PrestoBatchVectorSerializer.
TEST_P(PrestoSerializerTest, encodingsBatchVectorSerializer) {
    testBatchVectorSerializerRoundTrip(encodingsTestVector());
}

// Test that array elements have their encodings preserved by the
// PrestoBatchVectorSerializer.
TEST_P(PrestoSerializerTest, encodingsArrayElementsBatchVectorSerializer) {
    testBatchVectorSerializerRoundTrip(encodingsArrayElementsTestVector());
}

// Test that map values have their encodings preserved by the
// PrestoBatchVectorSerializer.
TEST_P(PrestoSerializerTest, encodingsMapValuesBatchVectorSerializer) {
    testBatchVectorSerializerRoundTrip(encodingsMapValuesTestVector());
}

// Test that dictionary encoding is preserved and that only the minimum number
// of entries from the alphabet are serialized.
TEST_P(PrestoSerializerTest, minimalDictionaryEncodings) {
    // Dictionary same size as rows.
    testMinimalDictionaryEncoding(32, 32);
    // Dictionary smaller than rows.
    testMinimalDictionaryEncoding(32, 16);
    // Dictionary larger than rows.
    testMinimalDictionaryEncoding(32, 64);
}

// Test that dictionary encoded inputs are flattened in cases where it doesn't
// help.
TEST_P(PrestoSerializerTest, dictionaryEncodingTurnedOff) {
    auto smallIntBase =
            make_flat_vector<int16_t>(32, [](vector_size_t row) { return row; });
    auto intBase =
            make_flat_vector<int32_t>(32, [](vector_size_t row) { return row; });
    auto bigintBase =
            make_flat_vector<int64_t>(32, [](vector_size_t row) { return row; });
    auto stringBase = make_flat_vector<std::string>(
        32, [](vector_size_t row) { return fmt::format("{}", row); });
    auto oneIndex = make_indices(32, [](auto) { return 0; });
    auto quarterIndices = make_indices(32, [](auto row) { return row % 8; });
    auto allButOneIndices = make_indices(32, [](auto row) { return row % 31; });
    auto allIndices = make_indices(32, [](auto row) { return row; });

    auto rows = make_row_vector({
        // Even though these have very effective dictionary encoding, they should
        // be flattened because their types are too small.
        BaseVector::wrap_in_dictionary(nullptr, oneIndex, 32, smallIntBase),
        BaseVector::wrap_in_dictionary(nullptr, oneIndex, 32, intBase),
        // These should keep dictionary encoding.
        BaseVector::wrap_in_dictionary(nullptr, oneIndex, 32, bigintBase),
        BaseVector::wrap_in_dictionary(nullptr, quarterIndices, 32, bigintBase),
        // These should be flattened because dictionary encoding isn't giving
        // enough benefit to outweigh the cost.
        BaseVector::wrap_in_dictionary(nullptr, allButOneIndices, 32, bigintBase),
        BaseVector::wrap_in_dictionary(nullptr, allIndices, 32, bigintBase),
        // These should keep dictionary encoding because strings are variable
        // length.
        BaseVector::wrap_in_dictionary(nullptr, oneIndex, 32, stringBase),
        BaseVector::wrap_in_dictionary(nullptr, quarterIndices, 32, stringBase),
        BaseVector::wrap_in_dictionary(nullptr, allButOneIndices, 32, stringBase),
        // This should be flattened because the alphabet is the same as the
        // flattened vector.
        BaseVector::wrap_in_dictionary(nullptr, allIndices, 32, stringBase),
    });

    for (bool preserveEncodings: {false, true}) {
        SCOPED_TRACE(fmt::format("preserveEncodings: {}", preserveEncodings));
        auto exptectedTransformedEncoding = preserveEncodings
                                                ? VectorEncoding::Simple::DICTIONARY
                                                : VectorEncoding::Simple::FLAT;
        serializer::presto::PrestoVectorSerde::PrestoOptions serdeOptions;
        serdeOptions.preserveEncodings = preserveEncodings;
        std::ostringstream out;
        serializeBatch(rows, &out, &serdeOptions);
        const auto serialized = out.str();

        auto rowType = as_row_type(rows->type());
        auto deserialized = deserialize(rowType, serialized, &serdeOptions);

        assertEqualVectors(rows, deserialized);

        // smallInt + one index
        ASSERT_EQ(
            deserialized->childAt(0)->encoding(), exptectedTransformedEncoding);
        // int + one index
        ASSERT_EQ(
            deserialized->childAt(1)->encoding(), exptectedTransformedEncoding);
        // bigint + one index
        ASSERT_EQ(
            deserialized->childAt(2)->encoding(),
            VectorEncoding::Simple::DICTIONARY);
        // bigint + quarter indices
        ASSERT_EQ(
            deserialized->childAt(3)->encoding(),
            VectorEncoding::Simple::DICTIONARY);
        // bigint + all but one indices
        ASSERT_EQ(
            deserialized->childAt(4)->encoding(), exptectedTransformedEncoding);
        // bigint + all indices
        ASSERT_EQ(
            deserialized->childAt(5)->encoding(), exptectedTransformedEncoding);
        // string + one index
        ASSERT_EQ(
            deserialized->childAt(6)->encoding(),
            VectorEncoding::Simple::DICTIONARY);
        // string + quarter indices
        ASSERT_EQ(
            deserialized->childAt(7)->encoding(),
            VectorEncoding::Simple::DICTIONARY);
        // string + all but one indices
        ASSERT_EQ(
            deserialized->childAt(8)->encoding(),
            VectorEncoding::Simple::DICTIONARY);
        // string + all indices
        ASSERT_EQ(
            deserialized->childAt(9)->encoding(), exptectedTransformedEncoding);
    }
}

TEST_P(PrestoSerializerTest, emptyVectorBatchVectorSerializer) {
    // Serialize an empty RowVector.
    auto row_vector = makeEmptyTestVector();

    std::ostringstream out;
    serializeBatch(row_vector, &out, nullptr);

    auto rowType = as_row_type(row_vector->type());
    auto deserialized = deserialize(rowType, out.str(), nullptr);
    assertEqualVectors(deserialized, row_vector);
}

TEST_P(PrestoSerializerTest, serializeNoRowsSelectedBatchVectorSerializer) {
    std::ostringstream out;
    kumo::pollux::serializer::presto::PrestoOutputStreamListener listener;
    OStreamOutputStream output(&out, &listener);
    auto paramOptions = getParamSerdeOptions(nullptr);
    auto serializer = serde_->createBatchSerializer(pool_.get(), &paramOptions);
    auto row_vector = make_row_vector(
        {"a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l"},
        {
            make_flat_vector<bool>({true, false, true, false}),
            make_flat_vector<int8_t>({1, 2, 3, 4}),
            make_flat_vector<int16_t>({5, 6, 7, 8}),
            make_flat_vector<int32_t>({9, 10, 11, 12}),
            make_flat_vector<int64_t>({13, 14, 15, 16}),
            make_flat_vector<float>({17.0, 18.0, 19.0, 20.0}),
            make_flat_vector<double>({21.0, 22.0, 23.0, 24.0}),
            make_flat_vector<StringView>({"25", "26", "27", "28"}),
            make_flat_vector<Timestamp>({{29, 30}, {32, 32}, {33, 34}, {35, 36}}),
            make_row_vector(
                {"1", "2"},
                {
                    make_flat_vector<StringView>({"37", "38", "39", "40"}),
                    make_flat_vector<int32_t>({41, 42, 43, 44})
                }),
            make_array_vector<int32_t>({{45, 46}, {47, 48}, {49, 50}, {51, 52}}),
            make_map_vector<StringView, int32_t>(
                {
                    {{"53", 54}, {"55", 56}},
                    {{"57", 58}, {"59", 60}},
                    {{"61", 62}, {"63", 64}},
                    {{"65", 66}, {"67", 68}}
                })
        });
    const IndexRange noRows{0, 0};

    serializer->serialize(row_vector, melon::Range(&noRows, 1), &output);

    auto expected = makeEmptyTestVector();
    auto rowType = as_row_type(expected->type());
    auto deserialized = deserialize(rowType, out.str(), nullptr);
    assertEqualVectors(deserialized, expected);
}

TEST_P(PrestoSerializerTest, lazy) {
    constexpr int kSize = 1000;
    auto row_vector = makeTestVector(kSize);
    auto lazyVector = std::make_shared<LazyVector>(
        pool_.get(),
        row_vector->type(),
        kSize,
        std::make_unique<SimpleVectorLoader>([&](auto) { return row_vector; }));
    testRoundTrip(lazyVector);
}

TEST_P(PrestoSerializerTest, ioBufRoundTrip) {
    VectorFuzzer::Options opts;
    opts.timestampPrecision =
            VectorFuzzer::Options::TimestampPrecision::kMilliSeconds;
    opts.nullRatio = 0.1;
    VectorFuzzer fuzzer(opts, pool_.get());

    const size_t numRounds = 200;

    for (size_t i = 0; i < numRounds; ++i) {
        auto rowType = fuzzer.randRowType();
        auto inputRowVector = fuzzer.fuzzInputRow(rowType);
        auto outputRowVector = IOBufToRowVector(
            rowVectorToIOBuf(inputRowVector, *pool_), rowType, *pool_);

        assertEqualVectors(inputRowVector, outputRowVector);
    }
}

TEST_P(PrestoSerializerTest, roundTrip) {
    VectorFuzzer::Options opts;
    opts.timestampPrecision =
            VectorFuzzer::Options::TimestampPrecision::kMilliSeconds;
    opts.nullRatio = 0.1;
    VectorFuzzer fuzzer(opts, pool_.get());
    VectorFuzzer::Options nonNullOpts;
    nonNullOpts.timestampPrecision =
            VectorFuzzer::Options::TimestampPrecision::kMilliSeconds;
    nonNullOpts.nullRatio = 0;
    VectorFuzzer nonNullFuzzer(nonNullOpts, pool_.get());

    const size_t numRounds = 100;

    for (size_t i = 0; i < numRounds; ++i) {
        auto rowType = fuzzer.randRowType();
        auto inputRowVector = (i % 2 == 0)
                                  ? fuzzer.fuzzInputRow(rowType)
                                  : nonNullFuzzer.fuzzInputRow(rowType);
        serializer::presto::PrestoVectorSerde::PrestoOptions prestoOpts;
        // Test every 2/4 with struct nulls first.
        prestoOpts.nullsFirst = i % 4 < 2;
        testRoundTrip(inputRowVector, &prestoOpts);
    }
}

TEST_P(PrestoSerializerTest, encodedRoundtrip) {
    VectorFuzzer::Options opts;
    opts.timestampPrecision =
            VectorFuzzer::Options::TimestampPrecision::kMilliSeconds;
    opts.nullRatio = 0.1;
    opts.dictionaryHasNulls = false;
    VectorFuzzer fuzzer(opts, pool_.get());

    const size_t numRounds = 200;

    for (size_t i = 0; i < numRounds; ++i) {
        auto rowType = fuzzer.randRowType();
        auto inputRowVector = fuzzer.fuzzInputRow(rowType);
        serializer::presto::PrestoVectorSerde::PrestoOptions serdeOpts;
        serdeOpts.nullsFirst = i % 2 == 0;
        testBatchVectorSerializerRoundTrip(inputRowVector, true, &serdeOpts);
    }
}

TEST_P(PrestoSerializerTest, opaqueBatchVectorSerializer) {
    OpaqueType::registerSerialization<Foo>(
        "Foo", Foo::serialize, Foo::deserialize);
    auto inputVector = make_flat_vector<std::shared_ptr<void> >(
        3,
        [](vector_size_t row) { return Foo::create(row + 10); },
        [](vector_size_t row) { return row == 1; },
        OPAQUE<Foo>());
    auto inputRowVector = make_row_vector({inputVector});

    std::ostringstream out;
    serializeBatch(inputRowVector, &out, nullptr);

    auto rowType = as_row_type(inputRowVector->type());
    auto deserialized = deserialize(rowType, out.str(), nullptr);
    assertEqualVectors(inputRowVector, deserialized);
}

TEST_P(PrestoSerializerTest, opaqueInteractiveVectorSerializer) {
    OpaqueType::registerSerialization<Foo>(
        "Foo", Foo::serialize, Foo::deserialize);
    auto inputVector = make_flat_vector<std::shared_ptr<void> >(
        3,
        [](vector_size_t row) { return Foo::create(row + 10); },
        [](vector_size_t row) { return row == 1; },
        OPAQUE<Foo>());
    auto inputRowVector = make_row_vector({inputVector});

    std::ostringstream out;
    POLLUX_ASSERT_THROW(
        serialize(inputRowVector, &out, nullptr),
        "Opaque type support is not implemented");
}

TEST_P(PrestoSerializerTest, encodedConcatenation) {
    // Slow test, run only for no compression.
    if (GetParam() != common::CompressionKind::CompressionKind_NONE) {
        return;
    }

    std::vector<TypePtr> types = {ROW({{"s0", VARCHAR()}}), VARCHAR()};
    VectorFuzzer::Options nonNullOpts{
        .nullRatio = 0, .dictionaryHasNulls = false
    };
    VectorFuzzer nonNullFuzzer(nonNullOpts, pool_.get());
    VectorFuzzer::Options nullOpts{.nullRatio = 1, .dictionaryHasNulls = false};
    VectorFuzzer nullFuzzer(nullOpts, pool_.get());
    VectorFuzzer::Options mixedOpts{
        .nullRatio = 0.5, .dictionaryHasNulls = false
    };
    VectorFuzzer mixedFuzzer(mixedOpts, pool_.get());

    for (const auto &type: types) {
        auto flatNull = nullFuzzer.fuzzFlat(type, 12);
        auto flatNonNull = nonNullFuzzer.fuzzFlat(type, 13);

        std::vector<VectorPtr> vectors = {
            nullFuzzer.fuzzConstant(type, 10),
            nonNullFuzzer.fuzzConstant(type, 11),
            flatNonNull,
            flatNull,
            nonNullFuzzer.fuzzDictionary(flatNonNull),
            nonNullFuzzer.fuzzDictionary(flatNull),
            nullFuzzer.fuzzDictionary(flatNull),
        };

        if (type->is_row()) {
            auto &rowType = type->as<TypeKind::ROW>();
            auto row = make_row_vector(
                {rowType.nameOf(0)},
                {nonNullFuzzer.fuzzConstant(rowType.childAt(0), 14)});
            row->set_null(2, true);
            row->set_null(10, true);
            vectors.push_back(row);
        }
        std::vector<std::vector<VectorPtr> > permutations;
        std::vector<VectorPtr> temp;
        makePermutations(vectors, 4, temp, permutations);
        for (auto i = 0; i < permutations.size(); ++i) {
            serializer::presto::PrestoVectorSerde::PrestoOptions opts;
            opts.nullsFirst = i % 2 == 0;

            testEncodedConcatenation(permutations[i], &opts);
        }
    }
}

TEST_P(PrestoSerializerTest, encodedConcatenation2) {
    // Slow test, run only for no compression.
    if (GetParam() != common::CompressionKind::CompressionKind_NONE) {
        return;
    }
    VectorFuzzer::Options options;
    VectorFuzzer fuzzer(options, pool_.get());
    for (auto nthType = 0; nthType < 20; ++nthType) {
        auto type = (fuzzer.randRowType());

        std::vector<VectorPtr> vectors;
        for (auto nth = 0; nth < 3; ++nth) {
            randomOptions(options, 0 == nth);
            fuzzer.setOptions(options);
            vectors.push_back(fuzzer.fuzzInputRow(type));
        }
        std::vector<std::vector<VectorPtr> > permutations;
        std::vector<VectorPtr> temp;
        makePermutations(vectors, 3, temp, permutations);
        for (auto i = 0; i < permutations.size(); ++i) {
            serializer::presto::PrestoVectorSerde::PrestoOptions opts;
            opts.useLosslessTimestamp = true;
            opts.nullsFirst = i % 2 == 0;

            testEncodedConcatenation(permutations[i], &opts);
        }
    }
}

TEST_P(PrestoSerializerTest, emptyArrayOfRowVector) {
    // The value of nullCount_ + nonNullCount_ of the inner RowVector is 0.
    auto arrayOfRow = make_array_of_row_vector(ROW({UNKNOWN()}), {{}});
    testRoundTrip(arrayOfRow);
}

TEST_P(PrestoSerializerTest, typeMismatch) {
    auto data = make_row_vector({
        make_flat_vector<int64_t>({1, 2, 3}),
        make_flat_vector<std::string>({"a", "b", "c"}),
    });

    std::ostringstream out;
    serialize(data, &out, nullptr);
    auto serialized = out.str();

    // Too many columns to deserialize.
    POLLUX_ASSERT_THROW(
        deserialize(ROW({BIGINT(), VARCHAR(), BOOLEAN()}), serialized, nullptr),
        "Number of columns in serialized data doesn't match "
        "number of columns requested for deserialization");

    // TMore columns in serialization than in type.
    if (GetParam() == common::CompressionKind_NONE) {
        // No throw.
        deserialize(ROW({BIGINT()}), serialized, nullptr);
    } else {
        POLLUX_ASSERT_THROW(
            deserialize(ROW({BIGINT()}), serialized, nullptr),
            "Number of columns in serialized data doesn't match "
            "number of columns requested for deserialization");
    }

    // Wrong types of columns.
    POLLUX_ASSERT_THROW(
        deserialize(ROW({BIGINT(), DOUBLE()}), serialized, nullptr),
        "Serialized encoding is not compatible with requested type: DOUBLE. "
        "Expected LONG_ARRAY. Got VARIABLE_WIDTH.");
}

TEST_P(PrestoSerializerTest, lexer) {
    VectorFuzzer::Options opts;
    opts.timestampPrecision =
            VectorFuzzer::Options::TimestampPrecision::kMilliSeconds;
    opts.nullRatio = 0.1;
    VectorFuzzer fuzzer(opts, pool_.get());
    VectorFuzzer::Options nonNullOpts;
    nonNullOpts.timestampPrecision =
            VectorFuzzer::Options::TimestampPrecision::kMilliSeconds;
    nonNullOpts.nullRatio = 0;
    VectorFuzzer nonNullFuzzer(nonNullOpts, pool_.get());

    const size_t numRounds = 20;

    for (size_t i = 0; i < numRounds; ++i) {
        auto rowType = fuzzer.randRowType();

        auto inputRowVector = (i % 2 == 0)
                                  ? fuzzer.fuzzInputRow(rowType)
                                  : nonNullFuzzer.fuzzInputRow(rowType);
        testLexer(inputRowVector);
    }
}

TEST_P(PrestoSerializerTest, checksum) {
    std::ostringstream output;
    // The payload (doesn't matter what as long as it's not 0).
    vector_size_t data = 4;

    // Write out the number of rows.
    vector_size_t numRows = std::numeric_limits<int32_t>::max();
    output.write(reinterpret_cast<char *>(&numRows), sizeof(numRows));
    // Set the bit indicating there's a checksum.
    char marker = 4;
    output.write(&marker, sizeof(marker));
    vector_size_t size = sizeof(data);
    // Write out the uncompressed size and size (we're not compressing the data so
    // they're the same).
    output.write(reinterpret_cast<char *>(&size), sizeof(size));
    output.write(reinterpret_cast<char *>(&size), sizeof(size));
    // Write out the checksum, it shouldn't match the checksum of data.
    int64_t checksum = 0;
    output.write(reinterpret_cast<char *>(&checksum), sizeof(checksum));
    output.write(reinterpret_cast<char *>(&data), sizeof(data));

    auto paramOptions = getParamSerdeOptions(nullptr);
    RowVectorPtr result;
    auto serialized = output.str();
    auto byteStream = toByteStream(serialized);
    // Make a small memory pool, if we try to allocate a Vector with numRows we'll
    // OOM.
    auto pool = memory::memoryManager()->addRootPool("checksum", 1UL << 10);
    // This should fail because the checksums don't match.
    POLLUX_ASSERT_THROW(
        serde_->deserialize(
            byteStream.get(),
            pool->addLeafChild("child").get(),
            ROW({BIGINT()}),
            &result,
            0,
            &paramOptions),
        "Received corrupted serialized page.");
}

INSTANTIATE_TEST_SUITE_P(
    PrestoSerializerTest,
    PrestoSerializerTest,
    ::testing::Values(
        common::CompressionKind::CompressionKind_NONE,
        common::CompressionKind::CompressionKind_ZLIB,
        common::CompressionKind::CompressionKind_SNAPPY,
        common::CompressionKind::CompressionKind_ZSTD,
        common::CompressionKind::CompressionKind_LZ4,
        common::CompressionKind::CompressionKind_GZIP));

TEST_F(PrestoSerializerTest, serdeSingleColumn) {
    // The difference between serialized data obtained from
    // PrestoIterativeVectorSerializer and serializeSingleColumn() is the
    // PrestoPage header and number of columns section in the serialized data.
    auto testSerializeRoundTrip = [&](const VectorPtr &vector) {
        static const size_t kPrestoPageHeaderBytes = 21;
        static const size_t kNumOfColumnsSerializedBytes = sizeof(int32_t);
        static const size_t kBytesToTrim =
                kPrestoPageHeaderBytes + kNumOfColumnsSerializedBytes;

        auto row_vector = make_row_vector({vector});
        // Serialize to PrestoPage format.
        std::ostringstream output;
        auto arena = std::make_unique<StreamArena>(pool_.get());
        auto rowType = as_row_type(row_vector->type());
        auto numRows = row_vector->size();
        auto serializer = serde_->createIterativeSerializer(
            rowType, numRows, arena.get(), nullptr);
        serializer->append(row_vector);
        kumo::pollux::serializer::presto::PrestoOutputStreamListener listener;
        OStreamOutputStream out(&output, &listener);
        serializer->flush(&out);

        validateLexer(output.str(), {});

        // Remove the PrestoPage header and Number of columns section from the
        // serialized data.
        std::string input = output.str().substr(kBytesToTrim);
        testDeserializeSingleColumn(input, vector);
    };

    auto testSerializeSingleColumnRoundTrip = [&](const VectorPtr &vector) {
        std::ostringstream output;
        serde_->serializeSingleColumn(vector, nullptr, pool_.get(), &output);
        const auto serialized = output.str();
        testDeserializeSingleColumn(serialized, vector);
    };

    // Verify that (de)serializeSingleColumn API can handle all supported types.
    std::vector<TypePtr> typesToTest = {
        BOOLEAN(),
        TINYINT(),
        SMALLINT(),
        INTEGER(),
        BIGINT(),
        REAL(),
        DOUBLE(),
        VARCHAR(),
        TIMESTAMP(),
        OPAQUE<Foo>(),
        ROW({VARCHAR(), INTEGER()}),
        ARRAY(INTEGER()),
        ARRAY(INTEGER()),
        MAP(VARCHAR(), INTEGER()),
        MAP(VARCHAR(), ARRAY(INTEGER())),
    };

    VectorFuzzer::Options opts;
    opts.vectorSize = 5;
    opts.nullRatio = 0.1;
    opts.dictionaryHasNulls = false;
    opts.stringVariableLength = true;
    opts.stringLength = 20;
    opts.containerVariableLength = false;
    opts.timestampPrecision =
            VectorFuzzer::Options::TimestampPrecision::kMilliSeconds;
    opts.containerLength = 10;

    auto seed = 0;

    KLOG(ERROR) << "Seed: " << seed;
    SCOPED_TRACE(fmt::format("seed: {}", seed));
    VectorFuzzer fuzzer(opts, pool_.get(), seed);
    fuzzer.registerOpaqueTypeGenerator<Foo>([](FuzzerGenerator &rng) {
        int64_t id = boost::random::uniform_int_distribution<int64_t>(1, 10)(rng);
        return Foo::create(id);
    });
    OpaqueType::registerSerialization<Foo>(
        "Foo", Foo::serialize, Foo::deserialize);

    for (const auto &type: typesToTest) {
        SCOPED_TRACE(fmt::format("Type: {}", type->toString()));
        auto data = fuzzer.fuzz(type);

        // Test deserializeSingleColumn() round trip with serialized data obtained
        // by PrestoIterativeVectorSerializer. This serialized data includes the
        // PrestoPage header and number of columns, which is removed for testing.
        testSerializeRoundTrip(data);

        // Test serializeSingleColumn() round trip with deserializeSingleColumn(),
        // both of these functions do not consider the PrestoPage header and number
        // of columns when (de)serializing the data.
        testSerializeSingleColumnRoundTrip(data);
    }
}

TEST_P(PrestoSerializerTest, nullVector) {
    std::ostringstream out;
    kumo::pollux::serializer::presto::PrestoOutputStreamListener listener;
    OStreamOutputStream output(&out, &listener);
    const auto paramOptions = getParamSerdeOptions(nullptr);
    const std::vector<IndexRange> ranges = {{0, 1}};
    Scratch scratch;
    auto serializer = serde_->createBatchSerializer(pool_.get(), &paramOptions);

    POLLUX_ASSERT_THROW(
        serializer->serialize(nullptr, ranges, scratch, &output),
        "Vector to serialize is null.");
}

TEST_P(PrestoSerializerTest, nullStream) {
    const auto paramOptions = getParamSerdeOptions(nullptr);
    const auto row =
            make_row_vector({BaseVector::create_null_constant(BIGINT(), 1, pool_.get())});
    auto serializer = serde_->createBatchSerializer(pool_.get(), &paramOptions);

    POLLUX_ASSERT_THROW(
        serializer->serialize(row, nullptr),
        "Stream to serialize out to is null.");
}

DEBUG_ONLY_TEST_P(PrestoSerializerTest, invalidRange) {
    std::ostringstream out;
    kumo::pollux::serializer::presto::PrestoOutputStreamListener listener;
    OStreamOutputStream output(&out, &listener);
    const auto paramOptions = getParamSerdeOptions(nullptr);
    const auto row =
            make_row_vector({BaseVector::create_null_constant(BIGINT(), 1, pool_.get())});
    Scratch scratch;
    auto serializer = serde_->createBatchSerializer(pool_.get(), &paramOptions);

    std::vector<IndexRange> ranges = {{-1, 1}};
    POLLUX_ASSERT_THROW(
        serializer->serialize(row, ranges, scratch, &output),
        "Invalid range at index 0");

    ranges = {{0, 2}};
    POLLUX_ASSERT_THROW(
        serializer->serialize(row, ranges, scratch, &output),
        "Invalid range at index 0");

    ranges = {{0, 1}, {-1, 1}};
    POLLUX_ASSERT_THROW(
        serializer->serialize(row, ranges, scratch, &output),
        "Invalid range at index 1");

    ranges = {{0, 1}, {1, 1}};
    POLLUX_ASSERT_THROW(
        serializer->serialize(row, ranges, scratch, &output),
        "Invalid range at index 1");
}

DEBUG_ONLY_TEST_P(PrestoSerializerTest, concurrency) {
    kumo::pollux::common::testutil::TestValue::enable();

    std::ostringstream out;
    kumo::pollux::serializer::presto::PrestoOutputStreamListener listener;
    OStreamOutputStream output(&out, &listener);
    const auto paramOptions = getParamSerdeOptions(nullptr);
    const auto row =
            make_row_vector({BaseVector::create_null_constant(BIGINT(), 1, pool_.get())});
    const std::vector<IndexRange> ranges = {{0, 1}};
    Scratch scratch;
    auto serializer = serde_->createBatchSerializer(pool_.get(), &paramOptions);

    melon::EventCount firstSerializeGo;
    std::atomic_bool firstSerializeGoFlag{false};
    melon::EventCount secondSerializeGo;
    std::atomic_bool secondSerializeGoFlag{false};

    SCOPED_TESTVALUE_SET(
        "kumo::pollux::serializers::PrestoBatchVectorSerializer::serialize",
        std::function<void(kumo::pollux::serializer::presto::detail::
            PrestoBatchVectorSerializer*)>(
            ([&](kumo::pollux::serializer::presto::detail::
                    PrestoBatchVectorSerializer* /*unused*/) {
                if (!secondSerializeGoFlag.load()) {
                secondSerializeGoFlag = true;
                secondSerializeGo.notifyAll();
                firstSerializeGo.await(
                    [&]() { return firstSerializeGoFlag.load(); });
                }
                })));

    std::thread t(
        [&]() { serializer->serialize(row, ranges, scratch, &output); });

    secondSerializeGo.await([&]() { return secondSerializeGoFlag.load(); });
    POLLUX_ASSERT_THROW(
        serializer->serialize(row, ranges, scratch, &output),
        "PrestoBatchVectorSerializer::serialize being called concurrently on the same object.");
    firstSerializeGoFlag = true;
    firstSerializeGo.notifyAll();

    t.join();
}

class PrestoSerializerBatchEstimateSizeTest : public testing::Test,
                                              public VectorBuilder {
protected:
    static void SetUpTestSuite() {
        if (!isRegisteredVectorSerde()) {
            serializer::presto::PrestoVectorSerde::registerVectorSerde();
        }
        ASSERT_EQ(getVectorSerde()->kind(), VectorSerde::Kind::kPresto);
        if (!isRegisteredNamedVectorSerde(VectorSerde::Kind::kPresto)) {
            serializer::presto::PrestoVectorSerde::registerNamedVectorSerde();
        }
        ASSERT_EQ(
            getNamedVectorSerde(VectorSerde::Kind::kPresto)->kind(),
            VectorSerde::Kind::kPresto);

        memory::MemoryManager::testingSetInstance({});
    }

    void SetUp() override {
        serde_ = std::make_unique<serializer::presto::PrestoVectorSerde>();
        serializer_ = serde_->createBatchSerializer(pool_.get(), &paramOptions_);
    }

    void testEstimateSerializedSize(
        VectorPtr input,
        const std::vector<IndexRange> &ranges,
        const std::vector<vector_size_t> &expectedSizes) {
        ASSERT_EQ(ranges.size(), expectedSizes.size());

        // Wrap the input a RowVector to better emulate production.
        auto row = make_row_vector({input});

        std::vector<vector_size_t> sizes(ranges.size(), 0);
        std::vector<vector_size_t *> sizesPtrs(ranges.size());
        for (int i = 0; i < ranges.size(); i++) {
            sizesPtrs[i] = &sizes[i];
        }

        Scratch scratch;
        serializer_->estimateSerializedSize(row, ranges, sizesPtrs.data(), scratch);

        for (int i = 0; i < expectedSizes.size(); i++) {
            // Add 4 bytes for each row in the wrapper. This is needed because we
            // wrap the input in a RowVector.
            ASSERT_EQ(sizes[i], expectedSizes[i] + 4 * ranges[i].size)
          << "Mismatched estimated size for range" << i << " "
          << ranges[i].begin << ":" << ranges[i].size;
        }
    }

    void testEstimateSerializedSize(
        const VectorPtr &vector,
        vector_size_t totalExpectedSize) {
        // The whole Vector is a single range.
        testEstimateSerializedSize(
            vector, {{0, vector->size()}}, {totalExpectedSize});
        // Split the Vector into two equal ranges.
        testEstimateSerializedSize(
            vector,
            {{0, vector->size() / 2}, {vector->size() / 2, vector->size() / 2}},
            {totalExpectedSize / 2, totalExpectedSize / 2});
        // Split the Vector into three ranges of 1/4, 1/2, 1/4.
        testEstimateSerializedSize(
            vector,
            {
                {0, vector->size() / 4},
                {vector->size() / 4, vector->size() / 2},
                {vector->size() * 3 / 4, vector->size() / 4}
            },
            {totalExpectedSize / 4, totalExpectedSize / 2, totalExpectedSize / 4});
    }

    std::unique_ptr<serializer::presto::PrestoVectorSerde> serde_;
    std::unique_ptr<BatchVectorSerializer> serializer_;
    serializer::presto::PrestoVectorSerde::PrestoOptions paramOptions_;
};

TEST_F(PrestoSerializerBatchEstimateSizeTest, flat) {
    auto flatBoolVector =
            make_flat_vector<bool>(32, [](vector_size_t row) { return row % 2 == 0; });

    // Bools are 1 byte each.
    // 32 * 1 = 32
    testEstimateSerializedSize(flatBoolVector, 32);

    auto flatIntVector =
            make_flat_vector<int32_t>(32, [](vector_size_t row) { return row; });

    // Ints are 4 bytes each.
    // 4 * 32 = 128
    testEstimateSerializedSize(flatIntVector, 128);

    auto flatDoubleVector =
            make_flat_vector<double>(32, [](vector_size_t row) { return row; });

    // Doubles are 8 bytes each.
    // 8 * 32 = 256
    testEstimateSerializedSize(flatDoubleVector, 256);

    auto flatStringVector = make_flat_vector<std::string>(
        32, [](vector_size_t row) { return fmt::format("{}", row); });

    // Strings are variable length, the first 10 are 1 byte each, the rest are 2
    // bytes.  Plus 4 bytes for the length of each string.
    // 10 * 1 + 22 * 2 + 4 * 32 = 182
    testEstimateSerializedSize(flatStringVector, {{0, 32}}, {182});
    testEstimateSerializedSize(flatStringVector, {{0, 16}, {16, 16}}, {86, 96});
    testEstimateSerializedSize(
        flatStringVector, {{0, 8}, {8, 16}, {24, 8}}, {40, 94, 48});

    auto flatVectorWithNulls = make_flat_vector<double>(
        32,
        [](vector_size_t row) { return row; },
        [](vector_size_t row) { return row % 4 == 0; });

    // Doubles are 8 bytes each, and only non-null doubles are counted. In
    // addition there's a null bit per row.
    // 8 * 24 + (32 / 8) = 196
    testEstimateSerializedSize(flatVectorWithNulls, {{0, 32}}, {196});
    testEstimateSerializedSize(
        flatVectorWithNulls, {{0, 16}, {16, 16}}, {98, 98});
    testEstimateSerializedSize(
        flatVectorWithNulls, {{0, 8}, {8, 16}, {24, 8}}, {49, 98, 49});
}

TEST_F(PrestoSerializerBatchEstimateSizeTest, array) {
    std::vector<vector_size_t> offsets{
        0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30
    };
    auto elements =
            make_flat_vector<int32_t>(32, [](vector_size_t row) { return row; });
    auto array_vector = make_array_vector(offsets, elements);

    // The ints in the array are 4 bytes each, and the array length is another 4
    // bytes per row.
    // 4 * 32 + 4 * 16 = 192
    testEstimateSerializedSize(array_vector, 192);

    std::vector<vector_size_t> offsetsWithEmptyOrNulls{
        0, 4, 4, 8, 8, 12, 12, 16, 16, 20, 20, 24, 24, 28, 28, 32
    };
    auto arrayVectorWithEmptyArrays =
            make_array_vector(offsetsWithEmptyOrNulls, elements);

    // The ints in the array are 4 bytes each, and the array length is another 4
    // bytes per row.
    // 4 * 32 + 4 * 16 = 192
    testEstimateSerializedSize(arrayVectorWithEmptyArrays, 192);

    std::vector<vector_size_t> nullOffsets{1, 3, 5, 7, 9, 11, 13, 15};
    auto arrayVectorWithNulls =
            make_array_vector(offsetsWithEmptyOrNulls, elements, nullOffsets);

    // The ints in the array are 4 bytes each, and the array length is another 4
    // bytes per non-null row, and 1 null bit per row.
    // 4 * 32 + 4 * 8 + 16 / 8 = 162
    testEstimateSerializedSize(arrayVectorWithNulls, {{0, 16}}, {162});
    testEstimateSerializedSize(arrayVectorWithNulls, {{0, 8}, {8, 8}}, {81, 81});
    testEstimateSerializedSize(
        arrayVectorWithNulls, {{0, 4}, {4, 8}, {8, 4}}, {41, 81, 41});
}

TEST_F(PrestoSerializerBatchEstimateSizeTest, map) {
    std::vector<vector_size_t> offsets{
        0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30
    };
    auto keys =
            make_flat_vector<int32_t>(32, [](vector_size_t row) { return row; });
    auto values =
            make_flat_vector<double>(32, [](vector_size_t row) { return row; });
    auto map_vector = make_map_vector(offsets, keys, values);

    // The ints in the map are 4 bytes each, the doubles are 8 bytes, and the
    // map length is another 4 bytes per row. 4 * 32 + 8 * 32 + 4 * 16 = 448
    testEstimateSerializedSize(map_vector, 448);

    std::vector<vector_size_t> offsetsWithEmptyOrNulls{
        0, 4, 4, 8, 8, 12, 12, 16, 16, 20, 20, 24, 24, 28, 28, 32
    };
    auto mapVectorWithEmptyMaps =
            make_map_vector(offsetsWithEmptyOrNulls, keys, values);

    // The ints in the map are 4 bytes each, the doubles are 8 bytes, and the
    // map length is another 4 bytes per row. 4 * 32 + 8 * 32 + 4 * 16 = 448
    testEstimateSerializedSize(mapVectorWithEmptyMaps, 448);

    std::vector<vector_size_t> nullOffsets{1, 3, 5, 7, 9, 11, 13, 15};
    auto mapVectorWithNulls =
            make_map_vector(offsetsWithEmptyOrNulls, keys, values, nullOffsets);

    // The ints in the map are 4 bytes each, the doubles are 8 bytes, and the
    // map length is another 4 bytes per non-null row, and 1 null bit per row.
    // 4 * 32 + 8 * 32 + 4 * 8 + 16 / 8 = 216
    testEstimateSerializedSize(mapVectorWithNulls, {{0, 16}}, {418});
    testEstimateSerializedSize(mapVectorWithNulls, {{0, 8}, {8, 8}}, {209, 209});
    testEstimateSerializedSize(
        mapVectorWithNulls, {{0, 4}, {4, 8}, {12, 4}}, {105, 209, 105});
}

TEST_F(PrestoSerializerBatchEstimateSizeTest, row) {
    auto field1 =
            make_flat_vector<int32_t>(32, [](vector_size_t row) { return row; });
    auto field2 =
            make_flat_vector<double>(32, [](vector_size_t row) { return row; });
    auto row_vector = make_row_vector({field1, field2});

    // The ints in the row are 4 bytes each, the doubles are 8 bytes, and the
    // offsets are 4 bytes per row.
    // 4 * 32 + 8 * 32 + 4 * 32 = 512
    testEstimateSerializedSize(row_vector, 512);

    auto rowVectorWithNulls =
            make_row_vector({field1, field2}, [](auto row) { return row % 4 == 0; });

    // The ints in the row are 4 bytes each, the doubles are 8 bytes, and the
    // offsets are 4 bytes per row, and 1 null bit per row.
    // 4 * 24 + 8 * 24 + 4 * 32 + 32 / 8 = 420
    testEstimateSerializedSize(rowVectorWithNulls, 420);
}

TEST_F(PrestoSerializerBatchEstimateSizeTest, constant) {
    auto flat_vector =
            make_flat_vector<int32_t>(32, [](vector_size_t row) { return row; });
    auto constantInt = BaseVector::wrap_in_constant(32, 0, flat_vector);

    // The single constant int is 4 bytes.
    testEstimateSerializedSize(constantInt, {{0, 32}}, {4});
    testEstimateSerializedSize(constantInt, {{0, 16}, {16, 16}}, {4, 4});
    testEstimateSerializedSize(
        constantInt, {{0, 8}, {8, 16}, {24, 8}}, {4, 4, 4});

    auto nullConstant = BaseVector::create_null_constant(BIGINT(), 32, pool_.get());

    // The single constant null is 1 byte (for the bit mask).
    testEstimateSerializedSize(nullConstant, {{0, 32}}, {1});
    testEstimateSerializedSize(nullConstant, {{0, 16}, {16, 16}}, {1, 1});
    testEstimateSerializedSize(
        nullConstant, {{0, 8}, {8, 16}, {24, 8}}, {1, 1, 1});

    std::vector<vector_size_t> offsets{
        0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30
    };
    auto array_vector = make_array_vector(offsets, flat_vector);
    auto constantArray = BaseVector::wrap_in_constant(32, 0, array_vector);

    // The single constant array is 4 bytes for the length, and 4 bytes for each
    // of the 2 integer elements.
    // 4 + 2 * 4 = 12
    testEstimateSerializedSize(constantArray, {{0, 32}}, {12});
    testEstimateSerializedSize(constantArray, {{0, 16}, {16, 16}}, {12, 12});
    testEstimateSerializedSize(
        constantArray, {{0, 8}, {8, 16}, {24, 8}}, {12, 12, 12});

    auto arrayVectorWithConstantElements = make_array_vector(offsets, constantInt);
    auto constantArrayWithConstantElements =
            BaseVector::wrap_in_constant(32, 0, arrayVectorWithConstantElements);

    // The single constant array is 4 bytes for the length, and 4 bytes for each
    // of the 2 integer elements (encodings for children of encoded complex
    // types are not currently preserved). 4 + 2 * 4 = 12
    testEstimateSerializedSize(
        constantArrayWithConstantElements, {{0, 32}}, {12});
    testEstimateSerializedSize(
        constantArrayWithConstantElements, {{0, 16}, {16, 16}}, {12, 12});
    testEstimateSerializedSize(
        constantArrayWithConstantElements,
        {{0, 8}, {8, 16}, {24, 8}},
        {12, 12, 12});
}

TEST_F(PrestoSerializerBatchEstimateSizeTest, dictionary) {
    auto indices = make_indices(32, [](auto row) { return (row * 2) % 32; });
    auto flat_vector =
            make_flat_vector<int32_t>(32, [](vector_size_t row) { return row; });
    auto dictionaryInts =
            BaseVector::wrap_in_dictionary(nullptr, indices, 32, flat_vector);

    // The indices are 4 bytes, and the dictionary entries are 4 bytes each.
    // 4 * 32 + 4 * 16 = 192
    testEstimateSerializedSize(dictionaryInts, {{0, 32}}, {192});
    testEstimateSerializedSize(dictionaryInts, {{0, 16}, {16, 16}}, {128, 128});
    testEstimateSerializedSize(
        dictionaryInts, {{0, 8}, {8, 16}, {24, 8}}, {64, 128, 64});

    auto flatVectorWithNulls = make_flat_vector<double>(
        32,
        [](vector_size_t row) { return row; },
        [](vector_size_t row) { return row % 4 == 0; });
    auto dictionaryNullElements =
            BaseVector::wrap_in_dictionary(nullptr, indices, 32, flatVectorWithNulls);

    // The indices are 4 bytes, half the dictionary entries are 8 byte doubles.
    // Note that the bytes for the null bits in the entries are not accounted
    // for, this is a limitation of having non-contiguous ranges selected from
    // the dictionary values. 4 * 32 + 8 * 8 = 192
    testEstimateSerializedSize(dictionaryNullElements, {{0, 32}}, {192});
    testEstimateSerializedSize(
        dictionaryNullElements, {{0, 16}, {16, 16}}, {128, 128});
    testEstimateSerializedSize(
        dictionaryNullElements, {{0, 8}, {8, 16}, {24, 8}}, {64, 128, 64});

    auto arrayIndices = make_indices(16, [](auto row) { return (row * 2) % 16; });
    std::vector<vector_size_t> offsets{
        0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30
    };
    auto array_vector = make_array_vector(offsets, flat_vector);
    auto dictionaryArray =
            BaseVector::wrap_in_dictionary(nullptr, arrayIndices, 16, array_vector);

    // The indices are 4 bytes, and the dictionary entries are 4 bytes length +
    // 4 bytes for each of the 2 array elements. 4 * 16 + 4 * 8 + 2 * 8 * 4 =
    // 160
    testEstimateSerializedSize(dictionaryArray, {{0, 16}}, {160});
    testEstimateSerializedSize(dictionaryArray, {{0, 8}, {8, 8}}, {128, 128});
    testEstimateSerializedSize(
        dictionaryArray, {{0, 4}, {4, 8}, {12, 4}}, {64, 128, 64});

    auto constantInt = BaseVector::wrap_in_constant(32, 0, flat_vector);
    auto arrayVectorWithConstantElements = make_array_vector(offsets, constantInt);
    auto dictionaryArrayWithConstantElements = BaseVector::wrap_in_dictionary(
        nullptr, arrayIndices, 16, arrayVectorWithConstantElements);

    // The indices are 4 bytes, and the dictionary entries are 4 bytes length +
    // 4 bytes for each of the 2 array elements (encodings for children of
    // encoded complex types are not currently preserved). 4 * 16 + 4 * 8 + 2 *
    // 8 * 4 = 160
    testEstimateSerializedSize(
        dictionaryArrayWithConstantElements, {{0, 16}}, {160});
    testEstimateSerializedSize(
        dictionaryArrayWithConstantElements, {{0, 8}, {8, 8}}, {128, 128});
    testEstimateSerializedSize(
        dictionaryArrayWithConstantElements,
        {{0, 4}, {4, 8}, {12, 4}},
        {64, 128, 64});

    auto dictionaryWithNulls = BaseVector::wrap_in_dictionary(
        make_nulls(32, [](auto row) { return row % 2 == 0; }),
        indices,
        32,
        flat_vector);

    // When nulls are present in the dictionary, currently we flatten the data.
    // So there are 4 bytes per row.  Null bits are only accounted for the null
    // elements because the non-null elements in the wrapped vector or
    // non-contiguous.
    // 4 * 16 + 16 / 8 = 66
    testEstimateSerializedSize(dictionaryWithNulls, {{0, 32}}, {66});
    testEstimateSerializedSize(
        dictionaryWithNulls, {{0, 16}, {16, 16}}, {33, 33});
    testEstimateSerializedSize(
        dictionaryWithNulls, {{0, 8}, {8, 16}, {24, 8}}, {17, 33, 17});
}
